#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>

#include "libtwclml.h"

/* 驱动设备路径 */
#define TWCLML_DEVICE_PATH "/dev/twclml"

/* 动态库全局状态 */
static int g_is_initialized = 0;  /* 初始化状态 */
static int g_dev_fd = -1;         /* 驱动设备文件描述符 */
static pthread_mutex_t g_mutex = PTHREAD_MUTEX_INITIALIZER;  /* 线程安全锁 */

/* 初始化驱动连接 */
nvmlReturn_t nvmlInit(void) {
    pthread_mutex_lock(&g_mutex);
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 已初始化直接返回 */
    if (g_is_initialized) {
        pthread_mutex_unlock(&g_mutex);
        return NVML_SUCCESS;
    }

    /* 打开驱动设备 */
    g_dev_fd = open(TWCLML_DEVICE_PATH, O_RDWR);
    if (g_dev_fd < 0) {
        pthread_mutex_unlock(&g_mutex);
        return NVML_ERROR_UNKNOWN;  /* 设备打开失败 */
    }

    g_is_initialized = 1;
    pthread_mutex_unlock(&g_mutex);
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlInit_v2(void) {
    return nvmlInit();  /* v2版本与v1一致 */
}

/* 关闭驱动连接 */
nvmlReturn_t nvmlShutdown(void) {
    pthread_mutex_lock(&g_mutex);
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 未初始化直接返回 */
    if (!g_is_initialized) {
        pthread_mutex_unlock(&g_mutex);
        return NVML_SUCCESS;
    }

    /* 关闭设备文件 */
    if (g_dev_fd != -1) {
        close(g_dev_fd);
        g_dev_fd = -1;
    }

    g_is_initialized = 0;
    pthread_mutex_unlock(&g_mutex);
    return NVML_SUCCESS;
}

/* 获取设备数量（调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetCount(unsigned int* count) {
    if (!g_is_initialized || !count) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_DEVICE_COUNT, count);
    pthread_mutex_unlock(&g_mutex);

    return (ret == 0) ? NVML_SUCCESS : NVML_ERROR_UNKNOWN;
}

nvmlReturn_t nvmlDeviceGetCount_v2(unsigned int* count) {
    return nvmlDeviceGetCount(count);  /* v2版本与v1一致 */
}

/* 通过索引获取设备句柄（调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device) {
    if (!g_is_initialized || !device) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 调用驱动IOCTL获取设备信息 */
    twclml_device_info_t drv_info;
    memset(&drv_info, 0, sizeof(drv_info));
    drv_info.index = index;

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_DEVICE_INFO, &drv_info);
    pthread_mutex_unlock(&g_mutex);

    if (ret != 0) {
        return NVML_ERROR_NOT_FOUND;  /* 设备不存在 */
    }

    /* 分配设备句柄并填充信息 */
    nvmlDevice_t dev = (nvmlDevice_t)malloc(sizeof(struct twcl_nvmlDevice));
    if (!dev) {
        return NVML_ERROR_UNKNOWN;  /* 内存分配失败 */
    }

    dev->index = drv_info.index;
    strncpy(dev->pciInfo.busId, drv_info.bus_id, sizeof(dev->pciInfo.busId)-1);
    dev->pciInfo.domain = drv_info.domain;
    dev->pciInfo.bus = drv_info.bus;
    dev->pciInfo.device = drv_info.device;
    dev->pciInfo.function = drv_info.function;
    dev->computeMajor = drv_info.compute_major;
    dev->computeMinor = drv_info.compute_minor;
    dev->numNvLinks = drv_info.num_nvlinks;

    *device = dev;
    return NVML_SUCCESS;
}

/* 获取设备PCI信息（从本地缓存读取） */
nvmlReturn_t nvmlDeviceGetPciInfo_v2(nvmlDevice_t device, nvmlPciInfo_t* pci) {
    if (!g_is_initialized || !device || !pci) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    *pci = device->pciInfo;  /* 直接从设备句柄缓存获取 */
    return NVML_SUCCESS;
}

/* 获取计算能力（从本地缓存读取） */
nvmlReturn_t nvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor) {
    if (!g_is_initialized || !device || !major || !minor) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    *major = device->computeMajor;  /* 从本地缓存获取 */
    *minor = device->computeMinor;
    return NVML_SUCCESS;
}

/* 获取NVLink状态（调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetNvLinkState(nvmlDevice_t device, unsigned int link, nvmlEnableState_t* isActive) {
    if (!g_is_initialized || !device || !isActive || link >= device->numNvLinks) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 调用驱动IOCTL获取NVLink状态 */
    twclml_nvlink_info_t drv_info;
    memset(&drv_info, 0, sizeof(drv_info));
    drv_info.device_index = device->index;
    drv_info.link = link;

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_NVLINK_INFO, &drv_info);
    pthread_mutex_unlock(&g_mutex);

    if (ret != 0) {
        return NVML_ERROR_UNKNOWN;
    }

    *isActive = drv_info.is_active ? NVML_ENABLED : NVML_DISABLED;
    return NVML_SUCCESS;
}

/* 获取NVLink远程PCI信息（调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned int link, nvmlPciInfo_t* pci) {
    if (!g_is_initialized || !device || !pci || link >= device->numNvLinks) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 调用驱动IOCTL获取远程PCI信息 */
    twclml_nvlink_info_t drv_info;
    memset(&drv_info, 0, sizeof(drv_info));
    drv_info.device_index = device->index;
    drv_info.link = link;

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_NVLINK_INFO, &drv_info);
    pthread_mutex_unlock(&g_mutex);

    if (ret != 0 || !drv_info.is_active) {
        return NVML_ERROR_NOT_SUPPORTED;
    }

    /* 填充远程PCI信息 */
    strncpy(pci->busId, drv_info.remote_bus_id, sizeof(pci->busId)-1);
    pci->domain = drv_info.remote_domain;
    pci->bus = drv_info.remote_bus;
    pci->device = drv_info.remote_device;
    pci->function = drv_info.remote_function;

    return NVML_SUCCESS;
}

/* 获取P2P状态（调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, 
                                   nvmlGpuP2PCapsIndex_t p2pIndex, 
                                   nvmlGpuP2PStatus_t* p2pStatus) {
    if (!g_is_initialized || !device1 || !device2 || !p2pStatus || p2pIndex >= NVML_P2P_CAPS_INDEX_UNKNOWN) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 调用驱动IOCTL获取P2P能力 */
    twclml_p2p_info_t drv_info;
    memset(&drv_info, 0, sizeof(drv_info));
    drv_info.device1_index = device1->index;
    drv_info.device2_index = device2->index;

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_P2P_INFO, &drv_info);
    pthread_mutex_unlock(&g_mutex);

    if (ret != 0) {
        return NVML_ERROR_UNKNOWN;
    }

    /* 根据能力索引返回对应状态 */
    switch (p2pIndex) {
        case NVML_P2P_CAPS_INDEX_READ:
            *p2pStatus = drv_info.read_supported ? NVML_P2P_STATUS_OK : NVML_P2P_STATUS_NOT_SUPPORTED;
            break;
        case NVML_P2P_CAPS_INDEX_WRITE:
            *p2pStatus = drv_info.write_supported ? NVML_P2P_STATUS_OK : NVML_P2P_STATUS_NOT_SUPPORTED;
            break;
        case NVML_P2P_CAPS_INDEX_NVLINK:
            *p2pStatus = drv_info.nvlink_supported ? NVML_P2P_STATUS_OK : NVML_P2P_STATUS_NOT_SUPPORTED;
            break;
        case NVML_P2P_CAPS_INDEX_ATOMICS:
            *p2pStatus = drv_info.atomics_supported ? NVML_P2P_STATUS_OK : NVML_P2P_STATUS_NOT_SUPPORTED;
            break;
        default:
            return NVML_ERROR_NOT_SUPPORTED;
    }

    return NVML_SUCCESS;
}

/* 错误信息字符串转换 */
const char* nvmlErrorString(nvmlReturn_t r) {
    switch (r) {
        case NVML_SUCCESS: return "Success";
        case NVML_ERROR_UNINITIALIZED: return "Library not initialized";
        case NVML_ERROR_INVALID_ARGUMENT: return "Invalid argument";
        case NVML_ERROR_NOT_SUPPORTED: return "Not supported";
        case NVML_ERROR_NOT_FOUND: return "Not found";
        default: return "Unknown error";
    }
}

/* 通过PCI Bus ID获取设备句柄（遍历设备匹配） */
nvmlReturn_t nvmlDeviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device) {
    if (!g_is_initialized || !pciBusId || !device) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    unsigned int count;
    nvmlReturn_t ret = nvmlDeviceGetCount(&count);
    if (ret != NVML_SUCCESS) {
        return ret;
    }

    /* 遍历所有设备匹配PCI Bus ID */
    for (unsigned int i = 0; i < count; i++) {
        nvmlDevice_t dev;
        ret = nvmlDeviceGetHandleByIndex(i, &dev);
        if (ret != NVML_SUCCESS) {
            continue;
        }

        if (strcmp(dev->pciInfo.busId, pciBusId) == 0) {
            *device = dev;
            return NVML_SUCCESS;
        }

        free(dev);  /* 不匹配的设备句柄释放 */
    }

    return NVML_ERROR_NOT_FOUND;  /* 未找到匹配设备 */
}

/* 获取设备索引（从本地缓存读取） */
nvmlReturn_t nvmlDeviceGetIndex(nvmlDevice_t device, unsigned int* index) {
    if (!g_is_initialized || !device || !index) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    *index = device->index;  /* 从本地缓存获取 */
    return NVML_SUCCESS;
}

/* 获取NVLink能力（完整实现，调用驱动IOCTL） */
nvmlReturn_t nvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, 
                                          nvmlNvLinkCapability_t capability, 
                                          unsigned int* capResult) {
    if (!g_is_initialized || !device || !capResult || link >= device->numNvLinks || capability >= NVML_NVLINK_CAP_COUNT) {
        return NVML_ERROR_INVALID_ARGUMENT;
    }
    printf("[%s +%d %s] \n", __FILE__, __LINE__, __func__);

    /* 调用驱动IOCTL获取NVLink能力 */
    twclml_nvlink_cap_info_t drv_info;
    memset(&drv_info, 0, sizeof(drv_info));
    drv_info.device_index = device->index;
    drv_info.link = link;
    drv_info.capability = capability;

    pthread_mutex_lock(&g_mutex);
    int ret = ioctl(g_dev_fd, TWCLML_IOCTL_GET_NVLINK_CAPABILITY, &drv_info);
    pthread_mutex_unlock(&g_mutex);

    if (ret != 0) {
        return NVML_ERROR_NOT_SUPPORTED;
    }

    *capResult = drv_info.cap_result;
    return NVML_SUCCESS;
}
    